from flask import json
from sqlalchemy import PrimaryKeyConstraint

from myapp.exts import db


class Consultation(db.Model):
    __tablename__ = 'consultations'

    userid = db.Column(db.Integer, nullable=False)
    consultationid = db.Column(db.Integer, nullable=False)
    content = db.Column(db.Text, nullable=True)

    # 定义复合主键
    __table_args__ = (PrimaryKeyConstraint('userid', 'consultationid'),)

    def __init__(self, userid, consultationid, content):
        self.userid = userid
        self.consultationid = consultationid
        self.content = content

    def to_dict(self):
        return {
            'userid': self.userid,
            'consultationid': self.consultationid,
            'content': self.content
        }

    def get_content_dict(self):
        content = self.content
        if content == '':
            messages = []
        else:
            messages = json.loads(content)
        return messages

    def add_answer(self, answer):
        messages = self.get_content_dict()
        msg = {
            'role': 'assistant',
            'content': answer
        }
        messages.append(msg)
        self.update_content(messages)

    def add_question(self, question):
        messages = self.get_content_dict()
        msg = {
            'role': 'user',
            'content': question
        }
        messages.append(msg)
        self.update_content(messages)

    def update_content(self, new_messages):
        self.content = json.dumps(new_messages, ensure_ascii=False)
        db.session.commit()
